import random

import numpy
import numpy as np

import matplotlib.pyplot as plt


class Node:
    weight = None
    phi = None
    id = None
    mu = 0.01
    c = 0.25
    a = 0.4
    target = None
    input = None
    target_wei = None

    def __init__(self, id, n: int, L: int, input, target_wei):
        self.weight = np.zeros(shape=(n, 1))
        self.id = id
        self.target = np.matmul(target_wei.T, input)
        self.target += np.random.randn(*self.target.shape) * 0.1
        self.input = input
        self.target_wei = target_wei

    def get_id(self):
        return self.id


class Graph:
    edges = {}
    nodes = {}
    mark_vis = {}

    def __init__(self):
        pass

    def add_node(self, node: Node):
        self.nodes[node.get_id()] = node
        self.edges[node.get_id()] = []

    def get_node_by_id(self, id) -> 'Node':
        return self.nodes[id]

    def add_edge_by_id(self, u, v):
        self.edges[u].append(v)

        self.edges[v].append(u)

    def get_agents_by_id(self, root):
        return self.edges[root]

    def iterate_edges(self):
        print("All edges are:")
        for k, v in self.edges.items():
            for e in v:
                print(f"{k} {e}")

    def clear_vis_tag(self):
        for k, v in self.nodes.items():
            self.mark_vis[k] = False

    def mark_visit_by_id(self, node_id):
        self.mark_vis[node_id] = True

    def visited_by_id(self, id):
        return self.mark_vis[id]


def incremental(graph: Graph, current_node: Node, T):
    graph.mark_visit_by_id(current_node.id)
    current_node.phi = current_node.weight
    for v in graph.get_agents_by_id(current_node.id):
        t_node = graph.get_node_by_id(v)
        current_node.phi = current_node.phi + current_node.mu * (
                t_node.c * np.expand_dims(t_node.input[:, T], axis=1) * (
                t_node.target[:, T] - np.matmul(np.expand_dims(t_node.input[:, T], axis=1).T, current_node.weight)
        )
        )
        if graph.visited_by_id(t_node.id):
            continue
        incremental(graph, t_node, T)


def diffusion(graph: Graph, current_node: Node, T):
    graph.mark_visit_by_id(current_node.id)
    current_node.weight = np.zeros_like(current_node.weight)
    for v in graph.get_agents_by_id(current_node.id):
        t_node = graph.get_node_by_id(v)
        current_node.weight = current_node.weight + current_node.a * t_node.phi
        if graph.visited_by_id(t_node.id):
            continue
        diffusion(graph, t_node, T)


class UnionFindSets:
    father = {}

    def __init__(self, ids: tuple or list):
        for x in ids:
            self.father[x] = x

    def find(self, x):
        if self.father[x] == x:
            return x
        self.father[x] = self.find(self.father[x])
        return self.father[x]

    def judge(self, x, y) -> 'bool':
        return self.find(x) == self.find(y)

    def union(self, x, y):
        self.father[self.find(x)] = self.find(y)


# for i in range(1, 10000):
#     func(mu=random.uniform(0, 2), a=random.uniform(0, 2), c=random.uniform(0, 2))
seed = 2022
random.seed(seed)
np.random.seed(seed)
n = 10
L = 1000
nod_cnt = 5
Nodes = []
node_ids = []
for i in range(0, nod_cnt):
    node_ids.append(i + 1)
    x = Node(i + 1, n=n, L=L, input=np.random.rand(n, L), target_wei=np.random.rand(n, 1))
    Nodes.append(x)

UFS = UnionFindSets(node_ids)
graph = Graph()

for node in Nodes:
    graph.add_node(node)

for i in range(1, nod_cnt):  # 并查集建树，保证联通
    u = random.randint(1, nod_cnt)
    v = random.randint(1, nod_cnt)
    while UFS.judge(u, v):
        u = random.randint(1, nod_cnt)
        v = random.randint(1, nod_cnt)
    graph.add_edge_by_id(u, v)
    UFS.union(u, v)

for i in range(1, random.randint(1, nod_cnt)):
    u = random.randint(1, nod_cnt)
    v = random.randint(1, nod_cnt)
    graph.add_edge_by_id(u, v)
graph.iterate_edges()  # 把图打印出来，方便CSAcademy查看

root = Nodes[0]  # 随便找一个节点作为根
Loss = []
for i in range(1, L + 1):
    graph.clear_vis_tag()
    incremental(graph=graph, current_node=root, T=i - 1)
    graph.clear_vis_tag()
    diffusion(graph=graph, current_node=root, T=i - 1)  # 更新参数
    # 统计loss
    loss = 0
    for node in Nodes:
        loss += np.linalg.norm(node.weight - node.target_wei)

    Loss.append(loss)

x = numpy.linspace(1, L, L)
plt.plot(x, Loss)
plt.xlabel("Time step")
plt.ylabel("error")
plt.show()
